"""
在微调阶段，数据的格式必须与我们在预训练部分使用的格式不同

使用了torch.utils中的Dataset类和BERT的分词器将数据转换为所需的格式

"""
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer
from data_read import train_data, val_data, test_data
from bert_config_movie import MAX_LEN, TRAIN_BATCH_SIZE, VALID_BATCH_SIZE

# 设置预训练的BERT模型名称  离线下载
model_name = "./bert-base-uncased/"
# 使用预训练的BERT模型的名称创建BERT分词器（Tokenizer）
tokenizer = BertTokenizer.from_pretrained(model_name, truncation=True, do_lower_case=True)


# 定义数据集类（QueryData），继承自PyTorch的Dataset类
class QueryData(Dataset):
    # 接收数据框（dataframe）、分词器（tokenizer）和最大长度（max_len）作为输入参数
    def __init__(self, dataframe, tokenizer, max_len):
        self.tokenizer = tokenizer
        # 查询的文本数据
        self.text = dataframe['query']
        # 标签数据
        self.targets = dataframe['label']
        self.max_len = max_len

    def __len__(self):
        # 返回数据集长度
        return len(self.text)

    # 根据索引获取一个样本的函数
    def __getitem__(self, index):
        text = str(self.text[index])
        # 将文本按空格拆分成词语，并使用空格重新连接，以确保所有的词语都以空格分隔
        text = " ".join(text.split())

        # 使用BERT分词器（tokenizer）将文本数据编码成BERT模型输入的格式
        inputs = self.tokenizer.encode_plus(
            text,
            None,
            add_special_tokens=True,
            max_length=self.max_len,
            pad_to_max_length=True,
            return_token_type_ids=True
        )

        ids = inputs['input_ids']
        # 获取注意力掩码（attention_mask），用于标识哪些词元是实际文本内容，哪些是填充的。
        mask = inputs['attention_mask']
        # 获取词元类型标识（token_type_ids），用于区分两个句子或段落
        token_type_ids = inputs["token_type_ids"]
        # 返回一个字典
        # BERT接受三个输入，即input_ids、attention_mask和token_type_ids
        return {
            # 将词元输入转换为PyTorch张量，并指定数据类型为torch.long（整型）
            'ids': torch.tensor(ids, dtype=torch.long),
            'mask': torch.tensor(mask, dtype=torch.long),
            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),
            # 将目标标签转换为PyTorch张量，并指定数据类型为torch.float（浮点型）
            'targets': torch.tensor(self.targets[index], dtype=torch.float)
        }


print("Train Dataset: {}".format(train_data.shape))
print("Validation Dataset: {}".format(val_data.shape))
print("Test Dataset: {}".format(test_data.shape))

training_set = QueryData(train_data, tokenizer, MAX_LEN)
val_set = QueryData(val_data, tokenizer, MAX_LEN)

# Defining training and testing paramerters
train_params = {'batch_size': TRAIN_BATCH_SIZE,
                'shuffle': True,
                'num_workers': 0
                }

test_params = {'batch_size': VALID_BATCH_SIZE,
               'shuffle': True,
               'num_workers': 0
               }

# Creating dataloader for training and testing purposes
training_loader = DataLoader(training_set, **train_params)
val_loader = DataLoader(val_set, **test_params)
